import torch
import torch.nn as nn
from model.nets.WindowTranformer_V2 import WindowTransformerV2
from model.nets.VisionTransformer import VisionAttention

class PixelShuffleDecoder(nn.Module):
    """ Pixel shuffle decoder. """

    # 对于HourGlass backbone，input_feat_dim=256
    def __init__(self, input_feat_dim=128, num_upsample=2, output_channel=2):
        super(PixelShuffleDecoder, self).__init__()
        # Get channel parameters
        self.channel_conf = self.get_channel_conf(num_upsample)  # [256, 64, 16]

        # Define the pixel shuffle
        self.pixshuffle = nn.PixelShuffle(2)  # 空间放大2倍，通道缩小为原来的1/4

        # Process the feature
        self.conv_block_lst = []
        # The input block
        self.conv_block_lst.append(
            nn.Sequential(
                nn.Conv2d(input_feat_dim, self.channel_conf[0],  # 256
                          kernel_size=3, stride=1, padding=1),
                nn.BatchNorm2d(self.channel_conf[0]),  # 256
                nn.ReLU(inplace=True)
            ))

        # Intermediate block
        for channel in self.channel_conf[1:-1]:  # 64
            self.conv_block_lst.append(
                nn.Sequential(
                    nn.Conv2d(channel, channel, kernel_size=3,
                              stride=1, padding=1),
                    nn.BatchNorm2d(channel),
                    nn.ReLU(inplace=True)
                ))

        # Output block
        self.conv_block_lst.append(
            nn.Conv2d(self.channel_conf[-1], output_channel,  # 16 2
                      kernel_size=1, stride=1, padding=0)
        )
        self.conv_block_lst = nn.ModuleList(self.conv_block_lst)

    # Get num of channels based on number of upsampling.
    def get_channel_conf(self, num_upsample):
        if num_upsample == 2:
            return [256, 64, 16]
        elif num_upsample == 3:
            return [256, 64, 16, 4]

    def forward(self, input_features):
        # Iterate til output block
        out = input_features
        # [:-1]从第1项到倒数第二项
        for block in self.conv_block_lst[:-1]:
            out = block(out)
            out = self.pixshuffle(out)

        # Output layer
        out = self.conv_block_lst[-1](out)

        return out

class MultiScaleAttention(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv0_to_1 = nn.Sequential(
            nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True)
        )
        self.downsample1 = nn.Conv2d(64, 128, kernel_size=1, stride=1, padding=0)
        self.bottleneck1_0 = nn.Sequential(
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 64, kernel_size=1, stride=1, padding=0),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 128, kernel_size=1, stride=1, padding=0),
        )
        self.bottleneck1_1 = nn.Sequential(
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.Conv2d(128, 64, kernel_size=1, stride=1, padding=0),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 128, kernel_size=1, stride=1, padding=0),
        )

        self.maxpool = nn.MaxPool2d(2, stride=2)
        self.downsample2 = nn.Conv2d(128, 256, kernel_size=1, stride=1, padding=0)
        self.bottleneck2_0 = nn.Sequential(
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.Conv2d(128, 128, kernel_size=1, stride=1, padding=0),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.Conv2d(128, 256, kernel_size=1, stride=1, padding=0),
        )
        self.bottleneck2_1 = nn.Sequential(
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 128, kernel_size=1, stride=1, padding=0),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.Conv2d(128, 256, kernel_size=1, stride=1, padding=0),
        )
        self.bottleneck2_2 = nn.Sequential(
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 128, kernel_size=1, stride=1, padding=0),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.Conv2d(128, 256, kernel_size=1, stride=1, padding=0),
        )
        self.bottleneck2_3 = nn.Sequential(
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, 128, kernel_size=1, stride=1, padding=0),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.Conv2d(128, 256, kernel_size=1, stride=1, padding=0),
        )

        self.maxp3 = nn.Sequential(
            nn.MaxPool2d(kernel_size=3, stride=2, padding=1),
            nn.Conv2d(256, 512, kernel_size=1, stride=1, padding=0),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True)
        )
        self.conv3 = nn.Sequential(
            nn.Conv2d(256, 512, kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True)
        )
        self.WA = WindowTransformerV2(scale_ratio=1, in_chans=512,
                                      depths=[2], num_heads=[8], mlp_ratio=[2.], window_size=8)

        self.maxp4 = nn.Sequential(
            nn.MaxPool2d(kernel_size=3, stride=2, padding=1),
            nn.Conv2d(512, 1024, kernel_size=1, stride=1, padding=0),
            nn.BatchNorm2d(1024),
            nn.ReLU(inplace=True)
        )
        self.conv4 = nn.Sequential(
            nn.Conv2d(512, 1024, kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(1024),
            nn.ReLU(inplace=True)
        )
        self.ViA = VisionAttention(16, 1024,
                                   patch_size=16,
                                   depth=6, num_heads=8, mlp_ratio=2.0)

        self.ps3 = nn.PixelShuffle(2)
        self.up3_0 = nn.Sequential(
            nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True)
        )
        self.up3_1 = nn.Sequential(
            nn.Conv2d(768, 512, kernel_size=1, stride=1, padding=0),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True)
        )

        self.ps2 = nn.PixelShuffle(2)
        self.up2_0 = nn.Sequential(
            nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True)
        )
        self.conv2_to_2 = nn.Sequential(
            nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True)
        )
        self.up2_1 = nn.Sequential(
            nn.Conv2d(384, 256, kernel_size=1, stride=1, padding=0),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True)
        )

        self.ps1 = nn.PixelShuffle(2)
        self.up1_0 = nn.Sequential(
            nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True)
        )
        self.conv1_to_1 = nn.Sequential(
            nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True)
        )
        self.up1_2 = nn.Sequential(
            nn.Conv2d(256, 128, kernel_size=1, stride=1, padding=0),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True)
        )

        self.ps0 = nn.PixelShuffle(2)
        self.up0 = nn.Sequential(
            nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True)
        )


    def forward(self, input_images):
        # 1/2 特征图
        x_1_0 = self.conv0_to_1(input_images)
        x_1_1 = self.downsample1(x_1_0) + self.bottleneck1_0(x_1_0)
        x_1_2 = self.bottleneck1_1(x_1_1) + x_1_1

        # 1/4 特征图
        x_2_0 = self.maxpool(x_1_2)
        x_2_1 = self.downsample2(x_2_0) + self.bottleneck2_0(x_2_0)
        x_2_2 = self.bottleneck2_1(x_2_1) + x_2_1
        x_2_3 = self.bottleneck2_2(x_2_2) + x_2_2
        x_2_4 = self.bottleneck2_3(x_2_3) + x_2_3

        # 1/8 特征图
        x_3_0 = self.maxp3(x_2_4) + self.conv3(x_2_4)
        x_3_1 = self.WA(x_3_0)

        # 1/16 特征图
        x_4_0 = self.maxp4(x_3_1[-1]) + self.conv4(x_3_1[-1])
        x_4_1 = self.ViA(x_4_0)

        # 融合 1/8
        x_3_2 = torch.cat((self.up3_0(self.ps3(x_4_1)), x_3_1[-1]), dim=1)
        x_3_4 = self.up3_1(x_3_2)

        # 融合 1/4
        x_2_5 = torch.cat((self.up2_0(self.ps2(x_3_4)), self.conv2_to_2(x_2_4)), dim=1)
        x_2_6 = self.up2_1(x_2_5)

        # 融合 1/2
        x_1_3 = torch.cat((self.up1_0(self.ps1(x_2_6)), self.conv1_to_1(x_1_2)), dim=1)
        x_1_4 = self.up1_2(x_1_3)

        # 融合 1
        x_0_0 = self.ps0(x_1_4)
        x_0_1 = self.up0(x_0_0)

        return x_0_1


class heatmap_head(nn.Module):
    '''
    基于 MUltiScaleAttention 输出的 H x W x 64 的特征图 和 angle分支输出的 H x W x 1 的角度图
    '''

    def __init__(self):
        super().__init__()
        self.predict_head = nn.Sequential(
            nn.Conv2d(64, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 2, kernel_size=1),
            nn.BatchNorm2d(2),
            nn.ReLU(inplace=True),
        )
        self.combine = nn.Sequential(
            nn.Conv2d(2, 2, kernel_size=1),
            nn.BatchNorm2d(2),
            nn.ReLU(inplace=True),
            nn.Conv2d(2, 2, kernel_size=1),
            nn.ReLU(inplace=True),
        )

    def forward(self, x, y):
        z = self.predict_head(x)

        q = torch.zeros_like(z)
        q[:, :1, :, :] = z[:, :1, :, :] - y
        q[:, 1:, :, :] = z[:, 1:, :, :] + y
        m = self.combine(q)

        return m